import pandas as pd
import csv
import shutil
from tqdm import tqdm
import hashlib
import os

sampleLableInfPath = r"/mnt/mydisk1/chenyongwei/malware/BODMAS_dataset/BODMAS_downloads/bodmas_metadata (1).csv"
bodmasSHAFamblePath = r"/home/cyw/projects/function_sim_project/malware_bazaar_dataset/bodmasSHAFamLable.txt"
bodmasMD5FamblePath = r"/home/cyw/projects/function_sim_project/malware_bazaar_dataset/bodmasMD5FamLable.txt"
originPath = r"/mnt/mydisk1/chenyongwei/malware/BODMAS_dataset/altered"
dataArmedPath = r"/mnt/mydisk1/chenyongwei/malware/BODMAS_dataset/armed"
dataTransPath = r"/mnt/mydisk1/chenyongwei/malware/BODMAS_dataset/dataTrans"


def getBodmasSampleNameAndLable():
    """
        获得bodmas样本集的恶意软件名及其家族标签
    """
    sample2Lable = {}
    fam2SampleList = {}
    with open(sampleLableInfPath, encoding='utf-8') as file_obj:
        reader_obj = csv.reader(file_obj)  # 以列表形式读取
        skipHead = True
        for row in reader_obj:  # 以list形式返回一行
            if skipHead:
                skipHead = False
                continue
            sampleName = row[0]
            sampleLable = row[2]
            if sampleLable != "":
                sample2Lable[sampleName] = sampleLable
                if sampleLable not in fam2SampleList:
                    fam2SampleList[sampleLable] = [sampleName]
                else:
                    fam2SampleList[sampleLable].append(sampleName)
        print("共{}个恶意软件".format(len(sample2Lable)))
        print("共{}个恶意家族".format(len(fam2SampleList)))
        with open(bodmasSHAFamblePath, 'w') as file:
            for sample in sample2Lable:
                file.write(sample+"\t"+sample2Lable[sample]+"\n")
    return [sample2Lable, fam2SampleList]


def filterSamples(fam2SampleList):
    """
        样本初步选择：
            移除掉样本数过少的家族
            样本过多的家族只选100个
        并生成恢复pe的中间文件
    """
    newFam2SampleList = {}
    seleceSample = []
    for i in fam2SampleList:
        if len(fam2SampleList[i]) < 10:
            continue
        if len(fam2SampleList[i]) > 100:
            newFam2SampleList[i] = fam2SampleList[i][:100]
        else:
            newFam2SampleList[i] = fam2SampleList[i]
    for i in newFam2SampleList:
        seleceSample.extend(newFam2SampleList[i])
    print("共选中{}个样本，{}个家族".format(len(seleceSample), len(newFam2SampleList)))

    # 获得所有待处理文件的绝对路径
    with open("/home/cyw/projects/function_sim_project/malware_bazaar_dataset/arm.txt", 'w') as file:
        for famLable in tqdm(newFam2SampleList):
            for sample in newFam2SampleList[famLable]:
                file.write(originPath+"//"+sample+".exe\n")
    return newFam2SampleList


def trans2Md5(newFam2SampleList):
    """
        处理样本，
            换成md5值和md52Fam
    """
    with open(bodmasMD5FamblePath, 'w') as file:
        for famLable in tqdm(newFam2SampleList):
            for sample in newFam2SampleList[famLable]:
                with open(dataArmedPath+"//"+sample+".exe", 'rb') as fp:
                    data = fp.read()
                fileMd5 = hashlib.md5(data).hexdigest()
                file.write(fileMd5+"\t"+famLable+"\n")
                shutil.copy(dataArmedPath+"//"+sample+".exe",
                            dataTransPath+"//"+fileMd5)
    print("数据转移完成")


def trans5000():
    """
        临时移动，后续删除
    """
    files = os.listdir(dataTransPath)
    cnt = 0
    trans1 = r"/mnt/mydisk1/chenyongwei/malware/BODMAS_dataset/halfTrans1"
    trans2 = r"/mnt/mydisk1/chenyongwei/malware/BODMAS_dataset/halfTrans2"
    for file in tqdm(files):
        cnt += 1
        if cnt < 5000:
            shutil.copy(dataTransPath+"//"+file, trans1+"//"+file)
        else:
            shutil.copy(dataTransPath+"//"+file, trans2+"//"+file)


if __name__ == "__main__":
    datas = getBodmasSampleNameAndLable()
    fam2SampleList = filterSamples(datas[1])
    # 这里需要使用 python PE_modifier.py --filelist arm.txt --altered_folder /mnt/mydisk1/chenyongwei/malware/BODMAS_dataset/armed --arm > restore.log
    # 才可以运行下一步
    # trans2Md5(fam2SampleList)

    # trans5000()
